﻿using System;
using System.Collections.Generic;
using System.Windows.Controls;
using System.Windows.Threading;

namespace DiskUsage
{
    /// <summary>
    /// http://bea.stollnitz.com/blog/?p=59
    /// How can I expand items in a TreeView? – Part III
    /// </summary>
    static class TreeViewExpander
    {
        public static void ExpandTreeViewItems<T>(Predicate<T> test, ItemsControl itemsControl)
        {
            ApplyActionToTreeViewItems(x =>
            {
                x.IsExpanded = true;
                DispatcherHelper.WaitForPriority(DispatcherPriority.ContextIdle);
            },
            test,
           itemsControl);
        }

        public static void SelectTreeViewItems<T>(Predicate<T> traverseTest,
            Predicate<TreeViewItem> equalsTest, ItemsControl itemsControl)
        {
            ApplyActionToTreeViewItems(x =>
            {
                if (equalsTest(x))
                {
                    x.IsSelected = true;
                    x.BringIntoView();
                }
            },
            traverseTest,
           itemsControl);
        }

        private static void ApplyActionToTreeViewItems<T>(Action<TreeViewItem> itemAction,
            Predicate<T> test, ItemsControl itemsControl)
        {
            Stack<ItemsControl> itemsControlStack = new Stack<ItemsControl>();
            itemsControlStack.Push(itemsControl);

            while (itemsControlStack.Count != 0)
            {
                ItemsControl currentItem = itemsControlStack.Pop() as ItemsControl;
                TreeViewItem currentTreeViewItem = currentItem as TreeViewItem;
                if (currentTreeViewItem != null)
                {
                    itemAction(currentTreeViewItem);
                }
                if (currentItem != null) //this handles the scenario where some TreeViewItems are already collapsed
                {
                    foreach (T dataItem in currentItem.Items)
                    {
                        if (test(dataItem))
                        {
                            ItemsControl childElement = (ItemsControl)currentItem.ItemContainerGenerator.
                                ContainerFromItem(dataItem);
                            itemsControlStack.Push(childElement);
                        }
                    }
                }
            }
        }
    }
}
